//! Python-oriented low-level IR that we can directly generate to file.

use std::string::ToString;

use thiserror::Error;

use stencil_tysys::{BinOp, IntEncoding, LiteralValue, MonOp, ValueType};
use stencil_util::{FieldName, FunctionName, IdentFormat, IdentString, ModuleName, TypeName};

#[derive(Debug, Error)]
pub enum Error {
    #[error("invalid identifier {0:?}")]
    InvalidIdentifier(String),
}

/// A python module making up a whole file.
pub struct PyModule {
    /// Comment inserted at top of file.
    pub initial_comment: Option<String>,

    /// The list of definitions.
    pub defs: Vec<PyDef>,
}

/// A definition or other item in a Python file.
pub enum PyDef {
    /// An import.
    Import(PyImport),

    /// A class definition.
    Class(PyClassDef),

    /// A function definition.
    Function(PyFnDef),

    /// A spurious comment.
    Comment(String),

    /// An guaranteed empty line.
    EmptyLine,
}

/// Imports of modules.
pub enum PyImport {
    /// A simple module import.
    Simple(ModuleSpec),

    /// Imports multiple defs from a module.
    Specific(ModuleSpec, Vec<Identifier>),
}

#[derive(Clone, Debug)]
pub struct ModuleSpec(String);

impl ModuleSpec {
    /// Constructs a module spec for a simple "single part" path, like for std.
    pub fn new_single(s: String) -> Result<Self, Error> {
        if is_str_identifier(&s) {
            Ok(Self(s))
        } else {
            Err(Error::InvalidIdentifier(s))
        }
    }

    /// Constructs a module from parts.
    pub fn from_parts<S: AsRef<str>, I: IntoIterator<Item = S>>(parts: I) -> Self {
        let mut s = String::new();

        // Annoying thing because .intersperse is nightly.
        let mut parts_iter = parts.into_iter();
        while let Some(part) = parts_iter.next() {
            let part_str = part.as_ref();
            if !is_str_identifier(part_str) {
                panic!("python/ir: invalid modulespec part '{part_str}'");
            }

            if !s.is_empty() {
                s.push_str(".");
            }

            s.push_str(part_str);
        }

        Self(s)
    }

    pub fn full(&self) -> &str {
        &self.0
    }

    pub fn parts(&self) -> impl Iterator<Item = &str> {
        self.0.split(".")
    }

    pub fn last(&self) -> &str {
        match self.0.rsplit_once(".") {
            Some((_, last)) => last,
            None => &self.0,
        }
    }

    pub fn into_string(self) -> String {
        self.0
    }
}

impl ToString for ModuleSpec {
    fn to_string(&self) -> String {
        self.0.clone()
    }
}

/// A formatted Python identifier.
#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Identifier(String);

impl Identifier {
    /// Converts from an ident string.
    pub fn from_ident_string(s: &IdentString, fmt: IdentFormat) -> Self {
        Self(s.to_format(fmt))
    }

    pub fn from_mod_name(s: &ModuleName) -> Self {
        Self::from_ident_string(s.inner(), IdentFormat::UpperCamel)
    }

    pub fn from_type_name(s: &TypeName) -> Self {
        Self::from_ident_string(s.inner(), IdentFormat::UpperCamel)
    }

    pub fn from_fn_name(s: &FunctionName) -> Self {
        Self::from_ident_string(s.inner(), IdentFormat::Snake)
    }

    pub fn from_field_name(s: &FieldName) -> Self {
        Self::from_ident_string(s.inner(), IdentFormat::Snake)
    }

    /// Tries to conver from a string.
    pub fn try_from(s: String) -> Result<Self, Error> {
        if is_str_identifier(&s) {
            Ok(Self(s))
        } else {
            Err(Error::InvalidIdentifier(s))
        }
    }

    pub fn into_string(self) -> String {
        self.0
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl AsRef<str> for Identifier {
    fn as_ref(&self) -> &str {
        &self.0
    }
}

/// A type as would be used in a type hint.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PyType {
    None,
    Str,
    Int,
    Float,
    Bool,

    /// Immutable bytes.
    Bytes,

    /// Mutable bytes.
    Bytearray,

    /// List of some type.
    List(Box<PyType>),

    /// Usually a user-owned class type.
    Class(Identifier),
}

impl PyType {
    pub fn from_mir(vt: &ValueType) -> Self {
        match vt {
            ValueType::Void => Self::None,

            ValueType::U8
            | ValueType::I8
            | ValueType::U16
            | ValueType::I16
            | ValueType::U32
            | ValueType::I32
            | ValueType::U64
            | ValueType::I64
            | ValueType::USize
            | ValueType::ISize => Self::Int,

            ValueType::Bool => Self::Bool,
            ValueType::ByteSlice => Self::Bytes,
            ValueType::StrUtf8Slice => Self::Str,
            ValueType::ByteArray(_) => Self::Bytearray,

            ValueType::ObjectInst(_) => {
                eprintln!("error (PyType::from_common_ir): don't support object types yet");
                unimplemented!()
            }

            ValueType::ObjectRef(_) => {
                eprintln!("error (PyType::from_common_ir): don't support object types yet");
                unimplemented!()
            }

            ValueType::Option(_) => {
                eprintln!("error (PyType::from_common_ir): don't support options yet");
                unimplemented!()
            }

            ValueType::List(ty) => Self::List(Box::new(Self::from_mir(ty))),
        }
    }

    pub fn to_python(&self) -> String {
        let s = match self {
            Self::None => "None",
            Self::Str => "str",
            Self::Int => "int",
            Self::Float => "float",
            Self::Bool => "bool",
            Self::Bytes => "bytes",
            Self::Bytearray => "bytearray",
            Self::List(inner) => return format!("list({})", inner.to_python()),
            Self::Class(ident) => return ident.as_str().to_owned(),
        };

        s.to_owned()
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PyValue {
    /// `None`
    None,

    /// A generic Python integer.
    // TODO support bigints properly
    Int(i64),

    /// A Python bool, ie `True` or `False`
    Bool(bool),
    // TODO more
}

impl PyValue {
    pub fn from_common_ir(lv: LiteralValue) -> Self {
        match lv {
            LiteralValue::Void => Self::None,
            LiteralValue::U8(v) => Self::Int(v as i64),
            LiteralValue::I64(v) => Self::Int(v),
            LiteralValue::U64(v) => Self::Int(v as i64),
            LiteralValue::USize(v) => Self::Int(v as i64),
            LiteralValue::ISize(v) => Self::Int(v),
            LiteralValue::Bool(v) => Self::Bool(v),
        }
    }

    /// Converts to a Python expression source code snippet.
    pub fn to_python(&self) -> String {
        match self {
            Self::None => "None".to_owned(),
            Self::Int(v) => format!("{v}"),
            Self::Bool(v) => (if *v { "True" } else { "False" }).to_owned(),
        }
    }
}

/// A class definition.
pub struct PyClassDef {
    /// The class name.
    pub name: Identifier,

    /// Documentation comment.
    pub doc: Option<String>,

    /// If the class should have `@dataclass` added to it.
    ///
    /// If this is set, should not have an explicit constructor.
    pub is_dataclass: bool,

    /// Class fields.
    pub fields: Vec<ClassFieldDef>,

    /// Local class definitions.
    pub functions: Vec<PyFnDef>,
}

/// A class field.
///
/// I am deciding that these must be typed.
pub struct ClassFieldDef {
    /// Name.
    pub name: Identifier,

    /// Type.
    pub ty: PyType,
}

impl ClassFieldDef {
    pub fn new(name: Identifier, ty: PyType) -> Self {
        Self { name, ty }
    }
}

/// The signature for a function, including names of each argument.
#[derive(Clone, Debug)]
pub struct FnSignature {
    /// If the function takes self.
    pub is_self: bool,

    /// Non-self arguments.
    pub args: Vec<FnArg>,

    /// Optional return type hint.
    pub ret_ty: Option<PyType>,
}

impl FnSignature {
    pub fn new(is_self: bool, args: Vec<FnArg>, ret_ty: Option<PyType>) -> Self {
        Self {
            is_self,
            args,
            ret_ty,
        }
    }

    pub fn empty() -> Self {
        Self::new(false, Vec::new(), None)
    }

    pub fn instance() -> Self {
        Self::new(true, Vec::new(), None)
    }

    pub fn with_arg(mut self, arg: FnArg) -> Self {
        self.args.push(arg);
        self
    }

    pub fn with_rty(mut self, ty: PyType) -> Self {
        self.ret_ty = Some(ty);
        self
    }
}

/// An argument to a function.
#[derive(Clone, Debug)]
pub struct FnArg {
    /// Name of the argument.
    pub name: Identifier,

    /// Type hint.
    pub ty: Option<PyType>,

    /// Default value
    ///
    /// Should match type hint.
    pub val: Option<PyValue>,
}

impl FnArg {
    pub fn new(name: Identifier, ty: Option<PyType>, val: Option<PyValue>) -> Self {
        Self { name, ty, val }
    }

    /// Creates a new instance with just a identifer for a name.
    pub fn new_simple(name: Identifier) -> Self {
        Self::new(name, None, None)
    }

    pub fn name(&self) -> &Identifier {
        &self.name
    }

    pub fn ty(&self) -> &Option<PyType> {
        &self.ty
    }

    pub fn val(&self) -> &Option<PyValue> {
        &self.val
    }

    /// Adds a type hint to the argument.
    pub fn with_type(mut self, ty: PyType) -> Self {
        self.ty = Some(ty);
        self
    }

    /// Adds a default value to the argument.
    pub fn with_val(mut self, val: PyValue) -> Self {
        self.val = Some(val);
        self
    }
}

#[derive(Clone, Debug)]
pub struct PyFnDef {
    /// The name of the function.  This may be a standalone fn or in a class.
    pub name: Identifier,

    /// Documentation comment.
    pub doc: Option<String>,

    /// If is the constructor.  I am not sure what this will be used for.
    pub is_constructor: bool,

    /// If within a class, if this should have `@staticmethod`.
    ///
    /// This is used for rich constructor fns.  Should not be set if not in a
    /// class.
    pub is_static: bool,

    /// The signature of the function.
    pub sig: FnSignature,

    /// The root block.
    pub body: PyBlock,
}

/// A block is a series of statements.
#[derive(Clone, Debug)]
pub struct PyBlock {
    pub statements: Vec<PyStatement>,
}

impl PyBlock {
    pub fn new(statements: Vec<PyStatement>) -> Self {
        Self { statements }
    }

    /// Creates a block of a single statement.
    pub fn single(statement: PyStatement) -> Self {
        Self::new(vec![statement])
    }
}

/// A statement is a procedural line within a function.
#[derive(Clone, Debug)]
pub enum PyStatement {
    /// A single-line comment.
    Comment(String),

    /// Assigns an expression to an lvalue.
    Assign(Lvalue, PyExpr),

    /// Assigns an expression to an index in a lvalue array.
    AssignIndex(Lvalue, PyExpr, PyExpr),

    /// Assignment into a buffer.
    ///
    /// (target, offset, value, enc)
    // TODO remove int type arg from this
    AssignBufInt(Lvalue, PyExpr, PyExpr, IntEncoding),

    /// Returns no value.
    ReturnNone,

    /// Returns a value.
    Return(PyExpr),

    /// Raises an exception.
    Raise(Exception),

    /// An if block.
    If(PyExpr, Box<PyBlock>),

    /// An if-else block.
    IfElse(PyExpr, Box<PyBlock>, Box<PyBlock>),

    /// A while loop with some condition.
    While(PyExpr, Box<PyBlock>),

    /// A for loop over an integer range.
    ForRange(Identifier, PyExpr, Box<PyBlock>),
}

/// An lvalue is a value that can be assigned to.
#[derive(Clone, Debug)]
pub enum Lvalue {
    /// A basic identifier reference.
    Ref(LvalueRef),

    /// A field on a local variable or argument.
    ///
    /// (ref, field)
    RefField(LvalueRef, Identifier),
}

impl Lvalue {
    pub fn new_self_field(id: Identifier) -> Self {
        Self::RefField(LvalueRef::ClassSelf, id)
    }
}

/// An owner for an lvalue.
#[derive(Clone, Debug)]
pub enum LvalueRef {
    /// The "self" reference on a class.
    ClassSelf,

    /// A local or a function arg.
    Local(Identifier),
}

/// A python expression.
#[derive(Clone, Debug)]
pub enum PyExpr {
    /// The "self" reference.
    ///
    /// Should only be used inside instance functions.
    ClassSelf,

    /// A literal value.
    Literal(PyValue),

    /// A local or a fn arg.
    ///
    /// Should not be used for "self" access, since there might be some special
    /// casing around that.
    Ident(Identifier),

    /// A ternary expression.
    ///
    /// (cond, if, else)
    Ternary(Box<PyExpr>, Box<PyExpr>, Box<PyExpr>),

    /// Cast to a type.
    ///
    /// This must be a "simple" type that can be cast to.
    Cast(PyType, Box<PyExpr>),

    /// Mono-op.
    Monop(MonOp, Box<PyExpr>),

    /// Binary-op.
    ///
    /// (op, lhs, rhs)
    Binop(BinOp, Box<PyExpr>, Box<PyExpr>),

    /// "Dot operator", for things like field access.
    Dot(Box<PyExpr>, Identifier),

    /// Index into a container.
    Index(Box<PyExpr>, Box<PyExpr>),

    /// Gets the length of a container.
    Len(Box<PyExpr>),

    /// Slice from a buffer.
    ///
    /// (target, start, end)
    BytebufReadSlice(Box<PyExpr>, Box<PyExpr>, Box<PyExpr>),

    /// Special-purpose integer slicing to extract an int from a buf.
    ///
    /// This would be annoying an awkward to express using general purpose
    /// expressions.
    ///
    /// This assumes that bounds checking is somehow done elsewhere.
    ///
    /// (target, offset, encoding)
    BytebufReadInt(Box<PyExpr>, Box<PyExpr>, IntEncoding),
}

impl PyExpr {
    /// Creates an expression for accessing a field on "self".
    pub fn new_self_field(id: Identifier) -> Self {
        Self::Dot(Box::new(Self::ClassSelf), id)
    }
}

/// An exception that can be raised.
#[derive(Clone, Debug)]
pub enum Exception {
    /// `ValueError`.
    ValueError(String),
}

/// Checks if the string is a safe identifier.  We only accept alphanum and "_".
fn is_str_identifier(s: &str) -> bool {
    if s.is_empty() {
        return false;
    }

    for c in s.chars() {
        if c.is_alphanumeric() {
            continue;
        }

        if c == '_' {
            continue;
        }

        return false;
    }

    true
}
